import cv2
import numpy as np

# Initialize variables
drawing = False  # True when the user is drawing
points = []  # List to store the polygon points
mask = None  # Binary mask for the selected region


def mouse_callback(event, x, y, flags, param):
    global drawing, points, mask

    if event == cv2.EVENT_LBUTTONDOWN:
        # Start drawing
        drawing = True
        points = [(x, y)]  # Initialize points with the first point

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            # Add points as the mouse moves
            points.append((x, y))
            # Draw the current polygon
            cv2.polylines(param, [np.array(points)], isClosed=False, color=(0, 255, 0), thickness=2)
            cv2.imshow('Image', param)

    elif event == cv2.EVENT_LBUTTONUP:
        # Finish drawing
        drawing = False
        # Close the polygon
        points.append(points[0])  # Close the loop
        cv2.polylines(param, [np.array(points)], isClosed=True, color=(0, 255, 0), thickness=2)
        cv2.imshow('Image', param)
        # Create a binary mask
        mask = np.zeros(param.shape[:2], dtype=np.uint8)
        cv2.fillPoly(mask, [np.array(points)], 255)  # Fill the polygon


def select_region(image):
    global mask, points
    mask = None  # Reset the mask
    points = []  # Reset the points

    # Clone the image
    img_copy = image.copy()

    # 直接使用imshow创建窗口，移除namedWindow调用
    cv2.imshow('Image', img_copy)
    cv2.setMouseCallback('Image', mouse_callback, param=img_copy)

    while True:
        cv2.imshow('Image', img_copy)
        key = cv2.waitKey(1) & 0xFF
        
        if key == 13 and mask is not None:  # Press Enter to confirm
            print("Selection confirmed.")
            break
        elif key == ord('q'):  # Press 'q' to abort
            print("Selection aborted.")
            mask = None
            break

    cv2.destroyAllWindows()
    return mask


if __name__ == "__main__":
    # Load an example image
    image = cv2.imread("example.jpg")
    if image is None:
        print("Error: Could not load the image.")
        exit()

    # Run the region selector
    region_mask = select_region(image)

    # Display the mask
    if region_mask is not None:
        cv2.imshow("Selected Region Mask", region_mask)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    else:
        print("No region was selected.")
